#include <cstdint>
#include <cstdio>
#include <map>
#include <string>
#include <utility>
#include <vector>

#include "opentelemetry/common/key_value_iterable_view.h"
#include "opentelemetry/sdk/trace/samplers/probability.h"
#include "opentelemetry/trace/span_context.h"
#include "opentelemetry/trace/span_context_kv_iterable_view.h"
#include "opentelemetry/trace/trace_state.h"
#include "src/common/random.h"

namespace trace_api = opentelemetry::trace;
using opentelemetry::sdk::common::Random;
using opentelemetry::sdk::trace::Decision;
using opentelemetry::sdk::trace::ProbabilitySampler;

namespace {
trace_api::TraceId TraceIdWithLow56(uint64_t v) {
  uint8_t buf[16] = {0};
  for (int i = 0; i < 7; ++i)
    buf[15 - i] = static_cast<uint8_t>(v >> (8 * i));
  return trace_api::TraceId(buf);
}

opentelemetry::sdk::trace::SamplingResult
Sample(ProbabilitySampler &s, const trace_api::SpanContext &ctx,
       trace_api::TraceId tid) {
  std::map<std::string, int> attrs;
  std::vector<
      std::pair<trace_api::SpanContext, std::map<std::string, std::string>>>
      links;
  opentelemetry::common::KeyValueIterableView<decltype(attrs)> attrs_view{
      attrs};
  trace_api::SpanContextKeyValueIterableView<decltype(links)> links_view{links};
  return s.ShouldSample(ctx, tid, "demo", trace_api::SpanKind::kInternal,
                        attrs_view, links_view);
}

const char *Header(const opentelemetry::sdk::trace::SamplingResult &r) {
  static std::string h;
  h = r.trace_state ? r.trace_state->ToHeader() : "(parent tracestate kept)";
  return h.c_str();
}
} // namespace

int main() {
  // ratio 1.0: always sampled, threshold 0 emitted as ot=th:0
  ProbabilitySampler all(1.0);
  auto r =
      Sample(all, trace_api::SpanContext::GetInvalid(), TraceIdWithLow56(0));
  printf("ratio 1.0    -> %s, tracestate: %s\n",
         r.decision == Decision::RECORD_AND_SAMPLE ? "SAMPLE" : "DROP",
         Header(r));

  // ratio 0.5: threshold is 2^55. R = 2^55 samples, R = 2^55-1 drops.
  ProbabilitySampler half(0.5);
  r = Sample(half, trace_api::SpanContext::GetInvalid(),
             TraceIdWithLow56(0x80000000000000));
  printf("ratio 0.5, R=0x80000000000000 -> %s, tracestate: %s\n",
         r.decision == Decision::RECORD_AND_SAMPLE ? "SAMPLE" : "DROP",
         Header(r));
  r = Sample(half, trace_api::SpanContext::GetInvalid(),
             TraceIdWithLow56(0x7fffffffffffff));
  printf("ratio 0.5, R=0x7fffffffffffff -> %s\n",
         r.decision == Decision::RECORD_AND_SAMPLE ? "SAMPLE" : "DROP");

  // 100000 random trace ids at ratio 0.5
  int sampled = 0;
  for (int i = 0; i < 100000; ++i) {
    uint8_t buf[16];
    Random::GenerateRandomBuffer(buf);
    if (Sample(half, trace_api::SpanContext::GetInvalid(),
               trace_api::TraceId(buf))
            .decision == Decision::RECORD_AND_SAMPLE)
      ++sampled;
  }
  printf("ratio 0.5, 100000 random trace ids -> sampled %d (%.3f)\n", sampled,
         sampled / 100000.0);

  // explicit rv in the parent tracestate wins over trace id randomness
  uint8_t id[16] = {1};
  trace_api::SpanContext rv_high(
      trace_api::TraceId(id), trace_api::SpanId(), trace_api::TraceFlags{0},
      false, trace_api::TraceState::FromHeader("ot=rv:ffffffffffffff"));
  r = Sample(half, rv_high, TraceIdWithLow56(0));
  printf("ratio 0.5, rv=ff.. trace id low bits 0 -> %s, tracestate: %s\n",
         r.decision == Decision::RECORD_AND_SAMPLE ? "SAMPLE" : "DROP",
         Header(r));

  // existing th is replaced, other entries survive
  ProbabilitySampler quarter(0.25);
  trace_api::SpanContext with_th(
      trace_api::TraceId(id), trace_api::SpanId(), trace_api::TraceFlags{0},
      false, trace_api::TraceState::FromHeader("ot=th:8,congo=t61rcWkgMzE"));
  r = Sample(quarter, with_th, TraceIdWithLow56(0xffffffffffffff));
  printf("ratio 0.25, parent 'ot=th:8,congo=..' -> %s, tracestate: %s\n",
         r.decision == Decision::RECORD_AND_SAMPLE ? "SAMPLE" : "DROP",
         Header(r));

  return 0;
}